import java.io.*;
import java.net.*;


public class TCPInfoExchangeServer {
    private static final int TCP_PORT = 4000;
    private static InetAddress clientAAddress;
    private static int clientAUDPPort;
    private static InetAddress clientBAddress;
    private static int clientBUDPPort;

    public static void main(String[] args) throws IOException {
        try (ServerSocket serverSocket = new ServerSocket(TCP_PORT)) {
            System.out.println("TCP Server is running on port " + TCP_PORT);

            // Accept connection from Client A
            Socket clientASocket = serverSocket.accept();
            exchangeInfo(clientASocket, true);

            // Accept connection from Client B
            Socket clientBSocket = serverSocket.accept();
            exchangeInfo(clientBSocket, false);

            // Start threads for UDP relay
            Thread threadAtoB = new Thread(() -> relayUDP(clientAAddress, clientAUDPPort, clientBAddress, clientBUDPPort));
            Thread threadBtoA = new Thread(() -> relayUDP(clientBAddress, clientBUDPPort, clientAAddress, clientAUDPPort));

            threadAtoB.start();
            threadBtoA.start();
        }
    }

    private static void exchangeInfo(Socket clientSocket, boolean isClientA) throws IOException {
        try (BufferedReader in = new BufferedReader(new InputStreamReader(clientSocket.getInputStream()));
             PrintWriter out = new PrintWriter(clientSocket.getOutputStream(), true)) {
            String clientMessage = in.readLine(); // Read client's UDP port
            if (isClientA) {
                clientAAddress = clientSocket.getInetAddress();
                clientAUDPPort = Integer.parseInt(clientMessage);
                System.out.println("Received from Client A: UDP port " + clientAUDPPort + " and address " + clientAAddress);
            } else {
                clientBAddress = clientSocket.getInetAddress();
                clientBUDPPort = Integer.parseInt(clientMessage);
                System.out.println("Received from Client B: UDP port " + clientBUDPPort + " and address " + clientBAddress);
            }
            out.println("Info received, you can start UDP communication.");
        }
    }

    private static void relayUDP(InetAddress sourceAddress, int sourcePort, InetAddress destinationAddress, int destinationPort) {
        try (DatagramSocket socket = new DatagramSocket(sourcePort)) {
            byte[] buffer = new byte[1024];
            while (true) {
                DatagramPacket packet = new DatagramPacket(buffer, buffer.length);
                socket.receive(packet);
                System.out.println("Received packet from " + sourcePort);

                packet.setAddress(destinationAddress);
                packet.setPort(destinationPort);
                socket.send(packet);
                System.out.println("Forwarded packet to " + destinationPort);
            }
        } catch (IOException e) {
            System.err.println("Error in UDP relay: " + e.getMessage());
        }
    }
}
